# 1. 环境与LLM初始化
import os
from dotenv import load_dotenv
from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory, ConversationSummaryMemory, \
    ConversationEntityMemory
from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0  # 0表示输出更确定，适合需要精确决策的场景
)

# 2. 定义工具（天气+计算器）
from langchain_core.tools import tool


@tool(description="查询指定城市的实时气温，参数为中文城市名（如'北京'），返回格式为'城市名：X℃'")
def WeatherTool(city: str) -> str:
    weather_data = {
        "北京": {"今天": "25℃", "明天": "23℃"},
        "上海": {"今天": "28℃", "明天": "26℃"},
        "广州": {"今天": "32℃", "明天": "30℃"}
    }
    return f"{city}今天：{weather_data.get(city, {}).get('今天', '暂无数据')}"


@tool(description="计算两个温度的差值，参数为两个温度字符串（用逗号分隔，如'28℃,25℃'），返回温差")
def TempDiffTool(temps: str) -> str:
    try:
        # 分割输入的两个温度
        temp1, temp2 = temps.split(',')
        t1 = int(temp1.replace("℃", "").strip())
        t2 = int(temp2.replace("℃", "").strip())
        return f"温差为{abs(t1 - t2)}℃"
    except:
        return "计算失败，请提供正确的温度格式（如'25℃,28℃'）"


tools = [WeatherTool, TempDiffTool]

# 3. 初始化记忆组件
# 核心:创建会话记忆
# memory = ConversationBufferMemory(
#     memory_key="chat_history",
#     return_messages=True,  # 返回Message对象，方便Agent理解
# )
# 提取实体（如城市、人名）并记忆
memory = ConversationEntityMemory(
    llm=llm,
    memory_key="chat_history",
    return_messages=True
)

# 4. 定义带工具的对话Prompt
from langchain_core.prompts import PromptTemplate

template = """
你是一个智能天气助手，能记住历史对话并调用工具回答问题。

历史对话：
{chat_history}

当前用户问题：{input}

可用工具：
{tools}

工具调用格式（如需使用）：
Action: 工具名
Action Input: 参数（如查询北京天气→北京）

注意：
1. 若用户问“明天呢”，需结合历史对话判断是哪个城市的明天；
2. 若用户问“温差”，需先获取两个城市的气温，再调用TempDiffTool；
3. 不需要工具时直接回答，回答要自然。
"""

prompt = PromptTemplate(
    template=ENTITY_MEMORY_CONVERSATION_TEMPLATE,
    input_variables=["input", "chat_history", "entities", "tools"],

)

prompt = PromptTemplate(
    template=ENTITY_MEMORY_CONVERSATION_TEMPLATE,
    input_variables=["input", "chat_history", "entities", "tools"]
)

# 5. 创建ConversationalAgent
from langchain.agents import initialize_agent, AgentType, AgentExecutor, ConversationalAgent

agent = ConversationalAgent.from_llm_and_tools(
    tools=tools,
    llm=llm,
    memory=memory,
    prompt=prompt,
    verbose=True
)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=10,
    memory=memory,  # 关键:把记忆绑定到执行器
)

# 6. 测试多轮对话+工具调用
if __name__ == "__main__":
    # 第1轮：查北京天气
    print("===== 第1轮 =====")
    print("Agent记忆1：", memory.buffer)
    res1 = agent_executor.invoke({"input": "北京今天气温多少？"})
    print("Agent回答：", res1["output"], "\n")

    # 第2轮：追问明天（依赖第1轮的“北京”）
    print("===== 第2轮 =====")
    print("Agent记忆2：", memory.buffer)
    res2 = agent_executor.invoke({"input": "那明天呢？"})
    print("Agent回答：", res2["output"], "\n")

    # 第3轮：查上海天气（新城市）
    print("===== 第3轮 =====")
    print("Agent记忆3：", memory.buffer)
    res3 = agent_executor.invoke({"input": "上海今天呢？"})
    print("Agent回答：", res3["output"], "\n")

    # 第4轮：算温差（依赖北京和上海的气温）
    print("===== 第4轮 =====")
    print("Agent记忆4：", memory.buffer)
    res4 = agent_executor.invoke({"input": "两地今天的温差是多少？"})
    print("Agent回答：", res4["output"])
